#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import sys
import time
import json
import struct
import shutil
import traceback
from pathlib import Path
from typing import Dict, Any, Optional, Tuple, Iterable
import threading                 # === 新增/保留：线程基元 ===
import queue                     # === 新增/保留：有界队列 ===
import torch
import argparse
# ===== 标准依赖（图像/数值） =====
import cv2
import numpy as np
from PIL import Image
from multiprocessing import Manager

# ===== 你原有的导入（保持） =====
from gears.netgear_udp import NetGearUDP as NetGear

# ===== 绑定地址与端口 =====
BIND_ADDR = "114.212.86.152"  # 127.0.0.1, 114.212.86.152
PORT = 5552

# ===== 回传给 client 的地址 =====
CLIENT_ADDR = "192.168.10.14"  # "172.27.144.86" "127.0.0.1""172.27.155.100"
CLIENT_PORT = 5553

# # ===== 保存目录（保持你的默认路径；首次运行清空重建） =====
# DATA_DIR = Path("/data/wxk/workspace/mirage/dataset/video000/Viduce/sender_out")
# SAVE_DIR = Path("/data/wxk/workspace/mirage/dataset/video000/Viduce/reciever_out")
# if SAVE_DIR.exists():
#     shutil.rmtree(SAVE_DIR)
# SAVE_DIR.mkdir(parents=True, exist_ok=True)


# ===== NS-3 收包函数（保持）=====
try:
    from ns3.receiver import receive_packet
except Exception as e:
    raise RuntimeError(f"[Init] 无法导入 ns3 函数 receive_packet: {e}")

# ===== ACK/RES 头部格式（保持）=====
_ACK_FMT = "!id"
_ACK_SIZE = struct.calcsize(_ACK_FMT)
RES_MAGIC = b'RSID'
RES_HDR_FMT = '!4sI'  # magic(4s) + frame_id(u32)



# --- 全局对象（带注释） ---
_q_decode = queue.Queue(maxsize=16)  # 有界队列：待解码的 frame_id；背压上游
_q_detect = queue.Queue(maxsize=16)  # 有界队列：待 YOLO 的 (frame_id, png_path, dec_ms)
_stop_event = threading.Event()      # 结束信号


def main():
    """
    主流程：
      - 接收 NS-3 送来的 UDP 包
      - receive_packet(data) → “刚刚完整的帧号列表”
      - 对每个新完成的帧 push -> _q_decode（阻塞时形成背压）
      - 两个工作线程常驻：_decoder_worker, _yolo_worker
      - 【新增】在此处一次性加载 YOLOContext，并传入 YOLO 工作线程
    """
    parser = argparse.ArgumentParser()
    parser.add_argument("--gcc", action="store_true",
                        help="enable GCC at receiver side (REMB/arrival-based control)")
    args = parser.parse_args()
    manager = Manager()
    done_frames_shared = manager.list()

    net = NetGear(
        address=BIND_ADDR,
        port=PORT,
        protocol="udp",
        receive_mode=True,
        logging=True,
        mtu=1500,
        recv_buffer_size=32 * 1024 * 1024,
        send_buffer_size=32 * 1024 * 1024,
        queue_maxlen=655360,
        slow_log_dir="./tmp/logs_receiver",          # 这端一般不会写 slow 日志，但可保留
        slow_weights_path="./tmp/slow_module_weights.json",
        rtcp_interval_ms=100,
    )
    net._peer_addr = (CLIENT_ADDR, CLIENT_PORT)
    net.enable_gcc(True)

    # === 在实例创建后，从实例上“绑定” pkt 常量到本地变量 ===
    # 这样后面代码仍可以用 PKT_DATA 等名字，不改流程。
    global PKT_DATA, PKT_SV_DATA, PKT_RES, PKT_TERM, PKT_ACK_FRAME, PKT_ACK_PACKET, PKT_PING, PKT_PONG
    PKT_DATA = net.PKT_DATA
    PKT_SV_DATA = net.PKT_SV_DATA
    PKT_RES = net.PKT_RES
    PKT_TERM = net.PKT_TERM
    PKT_ACK_FRAME = net.PKT_ACK_FRAME
    PKT_ACK_PACKET = net.PKT_ACK_PACKET
    PKT_PING = net.PKT_PING
    PKT_PONG = net.PKT_PONG

    seen = {}  # type: dict[int, float]

    # [新增变量] last_cleanup_ts：上次清理的时间戳，用于控制清理频率（每 10 秒一次）
    last_cleanup_ts = time.time()

    try:
        while True:
            pkt = net.recv()
            if pkt is not None:
                ptype = pkt.get("pkt_type")
                data = pkt.get("data", b"")

                if ptype in {PKT_DATA, PKT_SV_DATA, PKT_RES}:
                    # 解析收到的数据包，返回（本次刚刚变“完整”的帧ID列表, 当前包ID）
                    frame_ids, packet_id = receive_packet(data)

                    # 立即对“包”进行 ACK（包级 ACK）
                    if packet_id:
                        now = time.time()
                        ack_payload = struct.pack(_ACK_FMT, int(packet_id), now)
                        net.send(ack_payload, pkt_type=PKT_ACK_PACKET)

                    # 对“刚刚完整”的帧逐一去重并 ACK（帧级 ACK）
                    if frame_ids:
                        now = time.time()
                        for fid in frame_ids:
                            # 只对未见过的帧做 ACK，并记录时间戳
                            if fid not in seen:
                                seen[fid] = now  # 记录首次见到该帧的时间戳
                                ack_payload = struct.pack(_ACK_FMT, int(fid), now)
                                net.send(ack_payload, pkt_type=PKT_ACK_FRAME)

                # ===== 新增：每 10 秒清理一次 seen 中的过期项 =====
                # 功能：控制内存占用，把 10 秒前记录的 frame_id 移除（这些帧已完成 ACK、无需再保留）
                now = time.time()
                if now - last_cleanup_ts >= 10.0:
                    expire_before = now - 10.0
                    # 计算需要删除的键列表（避免边遍历边删除）
                    to_delete = [fid for fid, ts in seen.items() if ts < expire_before]
                    if to_delete:
                        for fid in to_delete:
                            seen.pop(fid, None)
                    last_cleanup_ts = now
                # ===== 新增清理结束 =====

                # 处理完当前 pkt，继续收下一个
                continue

            # 无包可收时小睡一会，降低 CPU 占用
            time.sleep(0.001)

    except KeyboardInterrupt:
        pass
    finally:
        _stop_event.set()
        # 尝试等队列清空（可选）
        try:
            _q_decode.join()
            _q_detect.join()
        except Exception:
            pass
        if hasattr(net, "get_stats"):
            print(f"[Server] Receiver stats: {net.get_stats()}")
        net.close()

if __name__ == "__main__":
    main()
